# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
from collections import namedtuple

from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.calculator import MovingSum
from coin.strategy.mm.simple_sim.executors import PassiveSimExecutor
from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.mm.subscription import get_subscription_key

FEATURE_ROW_FIELDS = [
    'time',
    'time_hr',
    'ask0_1',
    'askt_1',
    'bid0_1',
    'bidt_1',
    'ask0_2',
    'askt_2',
    'bid0_2',
    'bidt_2',
    'ask_basis_ma',
    'bid_basis_ma'
]
_FeatureRow = namedtuple('FeatureRow', FEATURE_ROW_FIELDS)


class FeatureRow(_FeatureRow):
  @staticmethod
  def get_csv_header():
    return 'header_feature_dump,%s' % ','.join(_FeatureRow._fields)

  def to_csv(self):
    return ('"FEATURE_DUMP",%d,"%s",'
            '%f,%f,%f,%f,%f,%f,%f,%f,%f,%f' % (self.time,
                                               self.time_hr,
                                               self.ask0_1,
                                               self.askt_1,
                                               self.bid0_1,
                                               self.bidt_1,
                                               self.ask0_2,
                                               self.askt_2,
                                               self.bid0_2,
                                               self.bidt_2,
                                               self.ask_basis_ma,
                                               self.bid_basis_ma))

  def to_log_string(self):
    return '\n' + self.to_csv()


class Pricer(object):
  def __init__(self,
               basis_ma_window,
               tick,
               book_askt_func_1,
               book_bidt_func_1,
               book_askt_func_2,
               book_bidt_func_2,
               price_multiplier_1=None):
    self._basis_ma_window = basis_ma_window
    self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick
    self._price_mult_1 = price_multiplier_1 or 1.
    self._book_askt_funcs = {1: book_askt_func_1, 2: book_askt_func_2}
    self._book_bidt_funcs = {1: book_bidt_func_1, 2: book_bidt_func_2}

    self._ready = False
    self._ask_basis_ma = MovingSum(self._basis_ma_window)
    self._bid_basis_ma = MovingSum(self._basis_ma_window)
    self._last = None

  @property
  def ready(self):
    return self._ready

  @property
  def last_value(self):
    return self._last

  def reset(self):
    self._ready = False
    self._last = None

  def update(self, book_1, book_2):
    book_ts = max(book_1.timestamp, book_2.timestamp)
    min_ts = min(book_1.timestamp, book_2.timestamp)

    if book_ts - min_ts > (2 * 10**9):
      self.reset()
      return

    if not (book_1.has_ask() and book_1.has_bid() and book_2.has_ask() and book_2.has_bid()):
      self.reset()
      return

    # Exchange 1
    ask0_1 = book_1.ask0().price
    bid0_1 = book_1.bid0().price
    askt_1 = self._book_askt_funcs[1](book_1)
    bidt_1 = self._book_bidt_funcs[1](book_1)

    # Exchange 2
    ask0_2 = book_2.ask0().price
    bid0_2 = book_2.bid0().price
    askt_2 = self._book_askt_funcs[2](book_2)
    bidt_2 = self._book_bidt_funcs[2](book_2)

    askt_diff = askt_1 * self._price_mult_1 - askt_2
    bidt_diff = bidt_1 * self._price_mult_1 - bidt_2

    self._ask_basis_ma.update(book_ts, askt_diff)
    self._bid_basis_ma.update(book_ts, bidt_diff)

    if len(self._ask_basis_ma) < 100 or len(self._bid_basis_ma) < 100:
      self.reset()
      return

    # Set ready
    self._last = FeatureRow(book_ts,
                            str(to_datetime(book_ts)),
                            ask0_1,
                            askt_1,
                            bid0_1,
                            bidt_1,
                            ask0_2,
                            askt_2,
                            bid0_2,
                            askt_2,
                            self._ask_basis_ma.average,
                            self._bid_basis_ma.average)
    self._ready = True

  def calculate_price(self, sell_edge, buy_edge, passive):
    assert self.ready
    assert -1 < sell_edge < 1
    assert -1 < buy_edge < 1

    askt_1_proj_2 = (self._last.askt_1 * self._price_mult_1 - float(self._ask_basis_ma.average))
    bidt_1_proj_2 = (self._last.bidt_1 * self._price_mult_1 - float(self._bid_basis_ma.average))

    if passive:
      ask0_minus_one_2 = self._tick.get_prev_price(self._last.ask0_2)
      bid0_plus_one_2 = self._tick.get_next_price(self._last.bid0_2)
      askp_2 = max(self._tick.ceil(askt_1_proj_2 * (1. + sell_edge)),
                   ask0_minus_one_2,
                   bid0_plus_one_2)
      bidp_2 = min(self._tick.floor(bidt_1_proj_2 * (1. - buy_edge)),
                   ask0_minus_one_2,
                   bid0_plus_one_2)
    else:
      askp_2 = self._tick.ceil(askt_1_proj_2 * (1. + sell_edge))
      bidp_2 = self._tick.floor(bidt_1_proj_2 * (1. - buy_edge))

    return askp_2, bidp_2


FILL_ROW_FIELDS = ['time', 'time_hr', 'side', 'price', 'qty', 'maker', 'pos', 'balance']
_FillRow = namedtuple('FillRow', FILL_ROW_FIELDS)


class FillRow(_FillRow):
  @staticmethod
  def get_csv_header():
    return 'header_fill_dump,%s' % ','.join(FILL_ROW_FIELDS)

  def to_csv(self):
    return ('"FILL_DUMP",%d,"%s",'
            '%d,%f,%f,%d,%f,%f' % (self.time,
                                   self.time_hr,
                                   self.side,
                                   self.price,
                                   self.qty,
                                   self.maker,
                                   self.pos,
                                   self.balance))

  def to_log_string(self):
    return '\n' + self.to_csv()


class PassUnhedgedSimStrategy(object):
  def __init__(
      self,
      pass_product,
      ref_product,
      sell_edge_func,  # pos --> sell_edge
      buy_edge_func,  # pos --> buy_edge
      pricing_param,
      executor_param,
      agg_edge=None,  # If None, does not aggress.
      trade_after=0,
      name=None,
      feature_filepath=None,
      fill_filepath=None):
    self._products = {1: ref_product, 2: pass_product}

    self._trade_after = to_timestamp_int(trade_after)
    self._start_ts = None

    self._sell_edge_func = sell_edge_func
    self._buy_edge_func = buy_edge_func
    self._agg_edge = agg_edge
    self._pricer = Pricer(**pricing_param)

    do_agg = agg_edge is not None

    unavail_intervals_file = None
    if pass_product.exchange == 'Bitmex':
      unavail_intervals_file = 'coin/strategy/data/bitmex_unavailability.txt'

    self._executor = PassiveSimExecutor(do_agg=do_agg,
                                        unavail_intervals_file=unavail_intervals_file,
                                        **executor_param)
    self._executor.pos_callback = self.on_pos_update
    self._books = {}

    self.name = name
    self._feature_filepath = feature_filepath
    self._feature_file = None
    self._last_feature_dump_ts = 0
    self._fill_filepath = fill_filepath
    self._fill_file = None
    self.fills = []

    self._last_ts = 0

  @property
  def book_ready(self):
    return (self._books.get(1, None) is not None and self._books.get(2, None) is not None)

  def on_book_reset(self, book_builder_name, book_builder):
    for product_id, product in self._products.items():
      subscription_key = get_subscription_key(product)
      # if not book_builder_name == subscription_key:
      #  continue

      book_builder.subscribe(product, functools.partial(self.on_book, product_id))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product_id))

  def on_book(self, product_id, book):
    if book.timestamp - self._last_ts < 10**8:
      return
    self._last_ts = book.timestamp

    if product_id == 2:
      self._executor.update_book(book)

    self._books[product_id] = book
    if not self.book_ready:
      return

    if self._start_ts is None:
      self._start_ts = book.timestamp + self._trade_after

    self._pricer.update(self._books[1], self._books[2])

    if not self._pricer.ready or book.timestamp < self._start_ts:
      return

    sell_edge = self._sell_edge_func(self._executor.pos.pos)
    buy_edge = self._buy_edge_func(self._executor.pos.pos)
    pass_sell_p, pass_buy_p = self._pricer.calculate_price(sell_edge, buy_edge, True)
    if self._agg_edge is not None:
      self._executor.cancel_agg_orders(book.timestamp)
      agg_sell_p, agg_buy_p = self._pricer.calculate_price(self._agg_edge, self._agg_edge, False)
      if agg_sell_p <= pass_sell_p:
        self._executor.agg_sell(agg_sell_p, book.timestamp)
      if agg_buy_p >= pass_buy_p:
        self._executor.agg_buy(agg_buy_p, book.timestamp)

    self._executor.post(book.timestamp, pass_sell_p, pass_buy_p)

    self._dump_feature()

  def on_trade(self, product_id, trade):
    if product_id == 2:
      self._executor.update_trade(trade)

  @property
  def pos(self):
    return self._executor.pos

  def get_summary(self):
    return self._executor.get_summary()

  def __del__(self):
    if self._feature_file is not None:
      self._feature_file.close()
    if self._fill_file is not None:
      self._fill_file.close()

  @property
  def _feature_dump_enabled(self):
    return self._feature_filepath is not None

  def _ensure_feature_file(self):
    if self._feature_dump_enabled and self._feature_file is None:
      self._feature_file = open(self._feature_filepath, 'w')
      self._feature_file.write(FeatureRow.get_csv_header() + '\n')

  def _dump_feature(self):
    if not self._feature_dump_enabled:
      return

    row = self._pricer.last_value
    if row is None or row.time - self._last_feature_dump_ts < 2 * (10**9):
      return
    self._last_feature_dump_ts = row.time

    self._ensure_feature_file()
    self._feature_file.write(row.to_csv() + '\n')

  @property
  def _fill_dump_enabled(self):
    return self._fill_filepath is not None

  def _ensure_fill_file(self):
    if self._fill_dump_enabled and self._fill_file is None:
      self._fill_file = open(self._fill_filepath, 'w')
      self._fill_file.write(FillRow.get_csv_header() + '\n')

  def on_pos_update(self, pos, price, qty, side, maker, timestamp):
    side = 1 if side == OrderSide.BUY else -1
    row = FillRow(timestamp, str(to_datetime(timestamp)), side, price, qty, maker, pos.pos, pos.pnl)
    self.fills.append(row)

    if not self._fill_dump_enabled:
      return

    self._ensure_fill_file()
    self._fill_file.write(row.to_csv() + '\n')
